package loveqq.jdbc.log;

import javassist.ClassPool;
import javassist.CtClass;
import loveqq.jdbc.log.config.AgentConfig;
import loveqq.jdbc.log.proxy.CoreProxy;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Properties;
import java.util.Set;

/**
 * JDBC 日志打印Log代理
 */
public class Agent {
    private static Set<String> driverProxyClass;

    static {
        try {
            // 读取driver.properties, 获取配置的Driver实现类
            Properties properties = new Properties();
            properties.load(Agent.class.getClassLoader().getResourceAsStream("driver.properties"));
            driverProxyClass = properties.stringPropertyNames();
        } catch (Exception e) {
            driverProxyClass = new HashSet<>();
        }
    }

    public static void premain(String agentArgs, Instrumentation instrumentation) throws Exception {
        // 入参配置
        AgentConfig config = AgentConfig.INSTANCE;
        config.load(agentArgs);
        if (config.isDebug()) {
            System.out.println("Agent配置: " + config);
        }

        /*
            插件思想：将JDBC整个过程的Connection，Statement全部代理，加上SQL+参数打印功能和一些辅助功能。
         */
        ClassPool pool = ClassPool.getDefault();
        CtClass preStmtClass = pool.get("java.sql.PreparedStatement");
        // 添加转换器
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (className == null) {
                    return classfileBuffer;
                }
                // 参数className是java/lang/String 转为 java.lang.String
                className = className.replaceAll("/", ".");
                // 需要代理的类
                CtClass ctClass = null;
                try {
                    // Driver实现类代理
                    if (driverProxyClass.contains(className)) {
                        ctClass = pool.get(className);
                        CoreProxy.proxyDriver(ctClass);
                        return ctClass.toBytecode();
                    }
                    // PreparedStatement实现类代理
                    if (className.contains("PreparedStatement")) {
                        ctClass = pool.get(className);
                        if (!ctClass.isInterface() && ctClass.subtypeOf(preStmtClass)) {
                            CoreProxy.proxyStatement(ctClass);
                            return ctClass.toBytecode();
                        }
                    }

                } catch (Throwable e) {
                    if (config.isDebug()) {
                        e.printStackTrace();
                    }
                } finally {
                    if (ctClass != null) {
                        ctClass.detach();
                    }
                }
                // 没有就不代理
                return classfileBuffer;
            }
        });
        // 释放资源
        preStmtClass.detach();
    }
}
